#####
# Replication for: "Can Political Speech Foster Tolerance of Immigrants?" by Schleiter, Tavits, and Ward.
# Figure S.1
#####

library(here)
library(data.table)
library(ggplot2)

# source the custom functions
source("functions.R")

# load the pooled data if not already in workspace
if(!exists("pooled")){
  pooled <- fread(here("data", "pooled.csv"))
}

# fit models --------------------------------------------------------------

#1. gender
bal_gend <- noCov(female)

#2. age
pooled[ , age_center := (age - mean(age, na.rm = T))/sd(age, na.rm = T), by = replication]
bal_age_center <- noCov(age_center)

#3. race
bal_white <- noCov(I(race2 == "White/Caucasian"))
bal_asian <- noCov(I(race2 == "Asian"))
bal_afam <- noCov(I(race2 == "Black/African American"))
bal_hisp <- noCov(I(race2 == "Hispanic/Latino"))
bal_other_race <- noCov(I(race2 == "Other (please specify)"))

#4. Education
bal_lessthan <- noCov(I(edu2 == "Less than high school"))
bal_highSchool <- noCov(I(edu2 == "High school graduate/ GED"))
bal_2year <- noCov(I(edu2 == "2 year degree"))
bal_someCollege <- noCov(I(edu2 == "Some college"))
bal_4year <- noCov(I(edu2 == "4 year degree"))
bal_masters <- noCov(I(edu2 == "Masters degree"))
bal_prof <- noCov(I(edu2 == "Professional Degree (JD, MD)"))
bal_doc <- noCov(I(edu2 == "Doctorate"))

#5. Party identification
bal_dem <- noCov(I(pid2 == "Democrat"))
bal_rep <- noCov(I(pid2 == "Republican"))
bal_ind <- noCov(I(pid2 == "Independent"))
bal_other_pid <- noCov(I(pid2 == "Something else"))

#6. News attention
bal_new <- noCov(news2)

#7. Region
bal_NE <- noCov(I(region2 == "Northeast"))
bal_MW <- noCov(I(region2 == "Midwest"))
bal_S <- noCov(I(region2 == "South"))
bal_W <- noCov(I(region2 == "West"))

#8. Wave
bal_w1 <- noCov(I(wave == "1")) 
bal_w2 <- noCov(I(wave == "2")) 
bal_w3 <- noCov(I(wave == "3")) 


# Prepare plotting --------------------------------------------------------

bal_models <- 
  list(
    bal_age_center,
    bal_gend,
    bal_white, bal_asian, bal_afam, bal_hisp, bal_other_race,
    bal_lessthan, bal_highSchool, bal_2year, bal_someCollege, bal_4year, bal_masters, bal_prof, bal_doc,
    bal_dem, bal_rep, bal_ind, bal_other_pid,
    bal_new,
    bal_NE, bal_MW, bal_S, bal_W,
    bal_w1, bal_w2, bal_w3
  )

bal_plot <- do.call("rbind", lapply(bal_models, function(mod){
  data.frame(extractPlotInfo(mod), 1, c("Common Humanity", "Countering Stereotypes","Norms"), stringsAsFactors = F)
}))

colnames(bal_plot) <- c("est", "lower95", "upper95", "lower90", "upper90", "x", "D")

bp_labs <-
  c("Age (centered to\nmean=0, SD=1):",
    "",
    "Female:",
    "",
    "Race:",
    "   White",
    "   Asian",
    "   African American/Black",
    "   Hispanic/Latino",
    "   Other",
    "",
    "Education:",
    "   Less than high school",
    "   High school grad./GED",
    "   2-year degree",
    "   Some college",
    "   4-year degree",
    "   Masters degree",
    "   Professional Degree",
    "   Doctorate",
    "",
    "Party Identification:",
    "   Democrat",
    "   Republican",
    "   Independent",
    "   Something else",
    "",
    "News Attention (1-7):",
    "",
    "Region:",
    "   Northeast",
    "   Midwest",
    "   South",
    "   West",
    "",
    "Survey Wave:",
    "   Wave 1",
    "   Wave 2", 
    "   Wave 3"
    ) 

bal_plot$x <- rev(rep(c(1:3, 6:9, 12, 14:17, 20:27, 30:34, 37, 39), each = 3))
bal_plot$D <- factor(bal_plot$D, levels = c("Norms", "Countering Stereotypes", "Common Humanity"))


# Plot and save -----------------------------------------------------------

fig_s1 <- ggplot(bal_plot, aes(x = est, y = x, xmin = lower95, xmax = upper95, group = D, color = D, shape = D)) +
  geom_vline(xintercept = 0, linetype = "dashed") +
  ggstance::geom_pointrangeh(fill = "white") +
  scale_y_continuous(breaks = length(bp_labs):1, labels = bp_labs, limits = c(.5, length(bp_labs) + 1), expand = expansion(0,0)) +
  ylab(NULL) +
  xlab("Regression Coefficients from Balance Tests") +
  scale_color_grey(name = "Treatment:", guide = guide_legend(reverse = T), start = .6, end = 0.2) +
  scale_shape_manual(values = c(17, 21, 15), name = "Treatment:", guide = guide_legend(reverse = T)) +
  theme(axis.text.y = element_text(hjust = 0, face = ifelse(length(bp_labs):1 %in% unique(bal_plot$x)[-c(1,2,20)] , "plain", "bold")),
        legend.position = "none",
        panel.grid.minor = element_blank()#,
        # axis.ticks.y = element_line(color = ifelse(bp_labs %in% c("", "Race:", "Education:", "Party Identification:", "Region:", "Survey Wave:"), "white","black")),
        # panel.grid.major.y = element_line(linetype = ifelse(bp_labs %in% c("", "Race:", "Education:", "Party Identification:", "Region:", "Survey Wave:"), "blank", "solid"))
  ) +
  facet_wrap( ~ D, ncol = 3)

ggsave(here("SI", "figure_s1.pdf"), fig_s1, width = 8, height = 8)
